
cmake_minimum_required(VERSION 3.16.0)

project(Ascend_C)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
# user-defined configuration
# set(SOC_VERSION "ascend910b4")
set(SOC_VERSION "ascend910b2")
set(DYNAMIC_MODE ON)
if(EXISTS /home/ma-user/Ascend/ascend-toolkit/latest)
  set(ASCEND_CANN_PACKAGE_PATH "/home/ma-user/Ascend/ascend-toolkit/latest")
elseif(EXISTS /usr/local/Ascend/ascend-toolkit/latest)
  set(ASCEND_CANN_PACKAGE_PATH "/usr/local/Ascend/ascend-toolkit/latest")
elseif(EXISTS /media/aaa/data/softs/Ascend/ascend-toolkit/latest)
  set(ASCEND_CANN_PACKAGE_PATH "/media/aaa/data/softs/Ascend/ascend-toolkit/latest")
endif()

message("cann地址：${ASCEND_CANN_PACKAGE_PATH}")

set(RUN_MODE "npu")
set(CMAKE_BUILD_TYPE "Release" CACHE STRING "Build type Release/Debug (default Debug)" FORCE)
set(CMAKE_INSTALL_PREFIX "${CMAKE_CURRENT_LIST_DIR}/out" CACHE STRING "path for install()" FORCE)

IF (CMAKE_HOST_SYSTEM_PROCESSOR MATCHES "aarch64")
    SET(ARCH aarch64)
ELSEIF (CMAKE_HOST_SYSTEM_PROCESSOR MATCHES "x86_64")
    SET(ARCH x86_64)
ENDIF ()

if(USE_CXX11_ABI)
    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=1 -fPIC -o2 -std=c++17")
else()
    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -o2 -std=c++17")
endif()

if(EXISTS ${ASCEND_CANN_PACKAGE_PATH}/tools/tikcpp/ascendc_kernel_cmake)
    set(ASCENDC_CMAKE_DIR ${ASCEND_CANN_PACKAGE_PATH}/tools/tikcpp/ascendc_kernel_cmake)
elseif(EXISTS ${ASCEND_CANN_PACKAGE_PATH}/compiler/tikcpp/ascendc_kernel_cmake)
    set(ASCENDC_CMAKE_DIR ${ASCEND_CANN_PACKAGE_PATH}/compiler/tikcpp/ascendc_kernel_cmake)
elseif(EXISTS ${ASCEND_CANN_PACKAGE_PATH}/ascendc_devkit/tikcpp/samples/cmake)
    set(ASCENDC_CMAKE_DIR ${ASCEND_CANN_PACKAGE_PATH}/ascendc_devkit/tikcpp/samples/cmake)
else()
    message(FATAL_ERROR "ascendc_kernel_cmake ${ASCEND_CANN_PACKAGE_PATH} does not exist, please check whether the cann package is installed.")
endif()

message("ASCENDC_CMAKE_DIR is ${ASCENDC_CMAKE_DIR}")
include(${ASCENDC_CMAKE_DIR}/ascendc.cmake)

ascendc_library(kernels STATIC add_custom/add_custom.cpp sub_custom/sub_custom.cpp)

ascendc_compile_definitions(kernels PRIVATE
  #$<$<BOOL:$<IN_LIST:${SOC_VERSION},${CUSTOM_ASCEND310P_LIST}>>:CUSTOM_ASCEND910B>
  #-DASCENDC_DUMP
  #-DHAVE_WORKSPACE
  #-DHAVE_TILING
  )

add_library(pybind11_lib SHARED pybind11.cpp)

target_link_libraries(pybind11_lib PRIVATE
    kernels
    torch_npu
)
execute_process(COMMAND python3 -c "import os; import torch; print(os.path.dirname(torch.__file__))"
  OUTPUT_STRIP_TRAILING_WHITESPACE
  OUTPUT_VARIABLE TORCH_PATH
)
message("TORCH_PATH is ${TORCH_PATH}")
set(ENV{ASCEND_HOME_PATH} ${ASCEND_CANN_PACKAGE_PATH})
execute_process(COMMAND python3 -c "import os; import torch_npu; print(os.path.dirname(torch_npu.__file__))"
  OUTPUT_STRIP_TRAILING_WHITESPACE
  OUTPUT_VARIABLE TORCH_NPU_PATH
)
message("TORCH_NPU_PATH is ${TORCH_NPU_PATH}")
target_link_directories(pybind11_lib PRIVATE
  ${TORCH_PATH}/lib
  ${TORCH_NPU_PATH}/lib
  ${ASCEND_CANN_PACKAGE_PATH}/lib64/
)

message("CMAKE_CURRENT_SOURCE_DIR is ${CMAKE_CURRENT_SOURCE_DIR}")
target_include_directories(pybind11_lib PRIVATE
  ${CMAKE_CURRENT_SOURCE_DIR}
  $<BUILD_INTERFACE:$<$<STREQUAL:${RUN_MODE},cpu>:${ASCEND_CANN_PACKAGE_PATH}/include>>
  $<BUILD_INTERFACE:$<$<STREQUAL:${RUN_MODE},cpu>:${ASCEND_CANN_PACKAGE_PATH}/runtime/include>>
  ${TORCH_NPU_PATH}/include
  ${TORCH_PATH}/include
  ${TORCH_PATH}/include/torch/csrc/api/include
  ${ASCEND_CANN_PACKAGE_PATH}/${ARCH}-linux/include/
  ${ASCEND_CANN_PACKAGE_PATH}/${ARCH}-linux/include/hccl/
)
include_directories(${CMAKE_CURRENT_SOURCE_DIR})

execute_process(COMMAND python3 -m pybind11 --includes
  OUTPUT_STRIP_TRAILING_WHITESPACE
  OUTPUT_VARIABLE PYBIND11_INC
)
string(REPLACE " " ";" PYBIND11_INC ${PYBIND11_INC})
target_compile_options(pybind11_lib PRIVATE
  ${PYBIND11_INC}
#  -D_GLIBCXX_USE_CXX11_ABI=0
#  -O2
#  -std=c++17
)

execute_process(COMMAND python3-config --extension-suffix
  OUTPUT_STRIP_TRAILING_WHITESPACE
  OUTPUT_VARIABLE PYBIND11_SUFFIX
)
set_target_properties(pybind11_lib PROPERTIES
  OUTPUT_NAME ascend_ops${PYBIND11_SUFFIX}
  PREFIX "" SUFFIX ""
)

install(TARGETS pybind11_lib DESTINATION ${CMAKE_INSTALL_PREFIX})
